import torch
import torch.nn as nn#nn是pytorch中torch.nn的别名，代表神经网络（neural network）模块
import torch.nn.functional as F


# Define the Resnet block，定义残差块
class ResnetBlock(nn.Module):#继承自nn.Module,是所有神经网络模块的基类，定义了神经网络的基本结构和方法
    def __init__(self, dim, use_dropout):#dim是输入的维度，use_dropout是是否使用dropout，use_dropout为True时，在残差块中添加dropout层
        #dropout:在训练过程中，随机丢弃一些神经元，防止过拟合
        #工作原理：在每次训练的前向传播时，Dropout 会以一定的概率（如 0.5）随机“丢弃”一部分神经元，即让它们的输出变为 0。这样，网络不会过度依赖某些特定的神经元，而是被迫让信息在网络中“分散”。
        super(ResnetBlock, self).__init__()#调用父类nn.Module的构造函数
        layers = [nn.ReflectionPad2d(1),#反射填充，在输入的周围填充1个像素，防止卷积操作时，边缘像素信息丢失
                  nn.Conv2d(in_channels=dim, out_channels=dim, kernel_size=3),#卷积层，输入维度为dim，输出维度为dim，卷积核大小为3
                  nn.InstanceNorm2d(dim),#实例归一化，将输入的维度归一化到0-1之间
                  nn.ReLU(inplace=True)]#激活函数，将输入的维度归一化到0-1之间
                    #ReLU:Rectified Linear Unit，线性整流单元，是一种常用的激活函数，定义为f(x)=max(0,x),即当x>0时，f(x)=x，当x<=0时，f(x)=0
                    #inplace=True:在原地进行操作，即不创建新的对象，而是直接修改原始对象。这样做可以节省内存，提高效率。
        if use_dropout:
            layers.append(nn.Dropout(0.5))#dropout层，将输入的维度随机丢弃50%
        
        layers += [nn.ReflectionPad2d(1),#反射填充，在输入的周围填充1个像素
                   nn.Conv2d(in_channels=dim, out_channels=dim, kernel_size=3),#卷积层，输入维度为dim，输出维度为dim，卷积核大小为3
                   nn.InstanceNorm2d(dim)]#实例归一化，将输入的维度归一化到0-1之间
        self.block = nn.Sequential(*layers)#将layers中的所有层组合成一个Sequential对象

    def forward(self, x):#定义前向传播，x是输入的维度  
        out = self.block(x) + x
        return out
    #这段代码实现了一个带有可选 Dropout 的标准残差块（Resnet Block），用于深度神经网络中，尤其适合图像生成和风格迁移任务。


# The Generator class，定义生成器
# The Generator consists of down-sampling layers, Resnet block, and up-sampling layers，生成器由下采样层、残差块和上采样层组成
class Generator(nn.Module):#继承自nn.Module,是所有神经网络模块的基类，定义了神经网络的基本结构和方法

    def __init__(self, in_channels, out_channels, out_features = 64, num_res_blocks = 9, use_dropout = False):
    #in_channels是输入的通道数，out_channels是输出的通道数，out_features是输出特征的维度，num_res_blocks是残差块的数量，use_dropout是是否使用dropout
        '''
        Parameters:
        in_channels: number of channels in input images
        out_channels: number of channels in output images
        out_features: number of channels after the first convolutional layer
        num_res_blocks: number of ResNet blocks
        use_dropput: whether to use dropout in the resnet block
        '''
        super(Generator, self).__init__()

        layers = [
            nn.ReflectionPad2d(3),#反射填充，在输入的周围填充3个像素，防止卷积操作时，边缘像素信息丢失
            nn.Conv2d(in_channels = in_channels, out_channels = out_features, kernel_size = 7),#7*7的卷积核，输入通道数为in_channels，输出通道数为out_features
            nn.InstanceNorm2d(out_features),#实例归一化，将输入的维度归一化到0-1之间
            nn.ReLU(inplace = True)#ReLU激活函数，将输入的维度归一化到0-1之间
        ]
        #第一层卷积层，输入通道数为in_channels，输出通道数为out_features，卷积核大小为7，填充为3，使用实例归一化，使用ReLU激活函数

        # Downsampling layers, there are 2 down-sampling layers by default，下采样层，默认有2个下采样层
        num_dowsampling = 2 #number of downsampling layers，下采样层的数量
        curr_channels = out_features#当前通道数
        for i in range(num_dowsampling):#循环num_dowsampling次
            next_channels = curr_channels * 2#下一层通道数
            layers.append(nn.Conv2d(in_channels = curr_channels, out_channels = next_channels, kernel_size = 3, stride = 2, padding = 1))
            #3*3的卷积核，输入通道数为curr_channels，输出通道数为next_channels，步长为2，填充为1
            layers.append(nn.InstanceNorm2d(next_channels))#实例归一化，将输入的维度归一化到0-1之间
            layers.append(nn.ReLU(inplace=True))#ReLU激活函数，将输入的维度归一化到0-1之间
            curr_channels = next_channels#更新当前通道数

        # Adding the resnet blocks，添加残差块
        for i in range(num_res_blocks):#循环num_res_blocks次
            layers.append(ResnetBlock(curr_channels, use_dropout = use_dropout))#添加残差块
            #堆叠残差块，将残差块的输出作为下一层的输入，增强网络的表达能力

        # Upsampling layers, there are 2 up-sampling layers by default，上采样层，默认有2个上采样层
        for i in range(num_dowsampling):
            next_channels = curr_channels // 2#下一层通道数为当前通道数的一半
            layers.append(nn.ConvTranspose2d(in_channels=curr_channels, out_channels=next_channels, kernel_size=3, stride=2, padding=1, output_padding=1))
            #ConvTranspose2d是转置卷积层/反卷积层/上采样层，3*3的卷积核，输入通道数为curr_channels，输出通道数为next_channels，步长为2，填充为1，输出填充为1
            layers.append(nn.InstanceNorm2d(next_channels))#实例归一化，将输入的维度归一化到0-1之间
            layers.append(nn.ReLU(inplace=True))#ReLU激活函数，将输入的维度归一化到0-1之间
            curr_channels = next_channels#更新当前通道数

        layers += [nn.ReflectionPad2d(3),#反射填充，在输入的周围填充3个像素，防止卷积操作时，边缘像素信息丢失
                   nn.Conv2d(in_channels = out_features, out_channels = out_channels, kernel_size = 7),
                   nn.Tanh()]#Tanh激活函数，将输入的维度归一化到-1到1之间

        self.model = nn.Sequential(*layers)
        '''
        这一步在 PyTorch 中非常常见，它的作用是将一组神经网络层（如卷积、归一化、激活等）按照顺序组合成一个整体模块，方便管理和调用。
        *layers 是 Python 中的解包操作符，用于将 layers 列表中的所有元素解包并传递给 nn.Sequential 构造函数。
        nn.Sequential 是 PyTorch 中用于按顺序组合多个层（如卷积、归一化、激活等）的模块。
        self.model 是 Generator 类的一个属性，用于存储整个生成器网络结构。
        '''
    def forward(self, x):#定义前向传播，x是输入的维度
        return self.model(x)
    #假设in_channels=3，out_channels=3，out_features=64，num_res_blocks=9，use_dropout=False，则layers的结构为：
    #1. 卷积层，经过一个7*7的卷积层，输出64个通道，然后经过一个InstanceNorm2d归一化，然后经过一个ReLU激活函数，输出64个通道
    #2. 下采样层，64个通道经过一个3*3的卷积层，输出128个通道，然后经过一个InstanceNorm2d归一化，然后经过一个ReLU激活函数，输出128个通道
    #3. 下采样层，128个通道经过一个3*3的卷积层，输出256个通道，然后经过一个InstanceNorm2d归一化，然后经过一个ReLU激活函数，输出256个通道
    #4. 残差块层，通道数保持不变
    #5. 上采样层，256个通道经过一个3*3的卷积层，输出128个通道，然后经过一个InstanceNorm2d归一化，然后经过一个ReLU激活函数，输出128个通道
    #6. 上采样层，128个通道经过一个3*3的卷积层，输出64个通道，然后经过一个InstanceNorm2d归一化，然后经过一个ReLU激活函数，输出64个通道
    #7. 输出层，64个通道经过一个7*7的卷积层，输出3个通道，然后经过一个Tanh激活函数，输出3个通道的图像
   


# The Discriminator class,定义判别器
class Discriminator(nn.Module):
    def __init__(self, in_channels, out_features = 64):
        '''
        Parameters
        in_channels: number of channels in input images
        out_features: number of channels after the first convolutional layer
        '''
        super(Discriminator, self).__init__()
        layers = []

        layers += [nn.Conv2d(in_channels=in_channels, out_channels=out_features, kernel_size=4, stride=2, padding=1),
                   nn.LeakyReLU(0.2, inplace=True)]
                    #LeakyReLU （Leaky Rectified Linear Unit，带泄漏的线性整流单元）是一种改进的 ReLU 激活函数，能让网络在负区间也有微弱的梯度，帮助模型更好地学习和收敛。
                    #其函数形式为：当x>0时，f(x)=x，当x<=0时，f(x)=a*x，其中a是一个很小的常数，此处取0.2。
                    #inplace=True:在原地进行操作，即不创建新的对象，而是直接修改原始对象。这样做可以节省内存，提高效率。

        curr_channels = out_features
        num_layers = 3 # number of convolutional layers

        for i in range(1, num_layers):
            next_channels = curr_channels * 2
            layers += [nn.Conv2d(in_channels= curr_channels, out_channels=next_channels, kernel_size=4, stride=2, padding=1),
                       nn.InstanceNorm2d(next_channels),
                       nn.LeakyReLU(0.2, inplace=True)]
            curr_channels = next_channels

        next_channels = curr_channels * 2
        layers += [nn.Conv2d(in_channels=curr_channels, out_channels=next_channels, kernel_size=4, stride=1, padding=1),
                   nn.InstanceNorm2d(next_channels),
                   nn.LeakyReLU(0.2, inplace=True)]
        curr_channels = next_channels

        # output one channel prediction map，输出一个通道的预测图，用于判断输入的图像是真实的还是生成的
        layers += [nn.Conv2d(in_channels= curr_channels, out_channels=1, kernel_size=4, stride=1, padding=1)]

        self.model = nn.Sequential(*layers)

    def forward(self, x):
        return self.model(x)
    #假设in_channels=3，out_features=64，则layers的结构为：
    #1. 输入图像经过一个4*4的卷积层，输出64个通道，然后经过一个LeakyReLU激活函数，输出64个通道
    #2. 64个通道经过一个4*4的卷积层，输出128个通道，然后经过一个InstanceNorm2d归一化，然后经过一个LeakyReLU激活函数，输出128个通道
    #3. 128个通道经过一个4*4的卷积层，输出256个通道，然后经过一个InstanceNorm2d归一化，然后经过一个LeakyReLU激活函数，输出256个通道
    #4. 256个通道经过一个4*4的卷积层，输出512个通道，然后经过一个InstanceNorm2d归一化，然后经过一个LeakyReLU激活函数，输出512个通道
    #5. 512个通道经过一个4*4的卷积层，输出1个通道，然后经过一个Tanh激活函数，输出一个通道的预测图，用于判断输入的图像是真实的还是生成的
'''
选择不同的卷积核会对模型的表达能力、参数量、感受野、训练效果等产生影响。
1.感受野(Receptive Field):
3x3卷积核感受野较小,但可以通过堆叠多层3x3卷积来获得更大的感受野。例如连续两层3x3卷积的感受野等价于一层5x5卷积。
4x4卷积核单层感受野更大,能一次性捕捉更大范围的特征。
结论:多层3x3卷积可以逐步扩大感受野,同时引入更多非线性,有助于提取更复杂的特征。
2.参数量(Parameter Count):
3x3卷积核:每层参数量为9。
4x4卷积核:每层参数量为16。
结论:4x4卷积核的参数量更少,但3x3卷积核可以堆叠更多层,获得更大的感受野。
3.非线性表达能力
3x3卷积核:堆叠多层3x3卷积,每层后面都可以加激活函数(如ReLU),引入更多非线性,有助于提升模型表达能力。
4x4卷积核:单层非线性较少。
结论:多层3x3卷积通常比单层大卷积核有更强的表达能力。
4.实际应用中的经验：
许多经典网络(如VGG、ResNet、UNet等)都偏好使用3x3卷积核,因为它们在参数效率和特征提取能力之间取得了很好的平衡。
4x4卷积核有时用于判别器(如GAN判别器),因为它可以更快地缩小空间尺寸,适合做PatchGAN等局部判别。
5.总结：
如果追求更强的特征表达能力、更深的网络结构,建议多用3x3卷积核。
如果需要快速缩小空间尺寸或做局部判别,可以用4x4卷积核。
在实际应用中,可以根据具体任务和数据集的特点选择合适的卷积核大小。
'''

#以下为U-net网络结构
# The building block of U-Net
# Two consecutive convolutional layers，双卷积块
class DoubleConv(nn.Module):
    def __init__(self, in_channels, out_channels):
        super(DoubleConv, self).__init__()
        layers = [
            nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1),
            nn.InstanceNorm2d(out_channels),
            nn.ReLU(inplace=True),
            nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1),
            nn.InstanceNorm2d(out_channels),
            nn.ReLU(inplace=True)
        ]
        self.layers = nn.Sequential(*layers)

    def forward(self, x):
        return self.layers(x)


# The building block of U-Net
# Down-sampling layers of U-Net，下采样块

class Down(nn.Module):
    def __init__(self, in_channels, out_channels):
        super(Down, self).__init__()
        layers = [
            nn.MaxPool2d(2),
            #MaxPool2d 是 PyTorch 中用于二维最大池化（Max Pooling）的一个层。
            #最大池化是一种常用的下采样（Downsampling）操作，通过在特征图上滑动一个固定大小的窗口，并选择窗口内的最大值作为输出。
            #它有助于减少特征图的尺寸，同时保留最重要的特征。
            #在U-Net中，最大池化通常用于下采样阶段，以减少特征图的尺寸，从而降低计算复杂度。
            DoubleConv(in_channels, out_channels)
        ]
        self.layers = nn.Sequential(*layers)

    def forward(self, x):
        return self.layers(x)
#先用2x2最大池化（MaxPool2d）将特征图尺寸减半（下采样），再用DoubleConv提取特征。

# The building block of U-Net
# Up-sampling layers of U-Net，上采样块

class Up(nn.Module):
    def __init__(self, in_channels, out_channels):
        super(Up, self).__init__()
        self.conv1 = nn.ConvTranspose2d(in_channels, in_channels // 2, kernel_size=2, stride=2)
        self.conv2 = DoubleConv(in_channels, out_channels)

    def forward(self, x1, x2):
        x1 = self.conv1(x1)
        x = torch.cat((x2, x1), dim=1)
        #这一步的作用是在通道维度（dim=1）上，将两个特征图拼接（concatenate）在一起。
        #torch.cat 是 PyTorch 中的张量拼接函数。
        #dim=1 表示在“通道”这个维度上拼接（对于形状为 [batch, channels, height, width] 的张量，dim=1 就是通道维度）。
        x = self.conv2(x)
        return x

#完整U-net网络结构
class U_net(nn.Module):
    def __init__(self, in_channels, out_channels, out_features = 64):
        super(U_net, self).__init__()
        self.conv1 = DoubleConv(in_channels, out_features)
        self.conv2 = Down(out_features, out_features * 2)
        self.conv3 = Down(out_features * 2, out_features * 4)
        self.conv4 = Down(out_features * 4, out_features * 8)
        self.conv5 = Down(out_features * 8, out_features * 16)
        self.deconv4 = Up(out_features * 16, out_features * 8)
        self.deconv3 = Up(out_features * 8, out_features * 4)
        self.deconv2 = Up(out_features * 4, out_features * 2)
        self.deconv1 = Up(out_features * 2, out_features)
        layers=  [
            nn.Conv2d(in_channels=out_features, out_channels=out_channels, kernel_size=3, padding=1),
            nn.Tanh()#Tanh激活函数，将输入的维度归一化到-1到1之间
        ]
        self.output = nn.Sequential(*layers)


    def forward(self, x):
        x1 = self.conv1(x)
        x2 = self.conv2(x1)
        x3 = self.conv3(x2)
        x4 = self.conv4(x3)
        x5 = self.conv5(x4)
        out = self.deconv4(x5, x4)
        out = self.deconv3(out, x3)
        out = self.deconv2(out, x2)
        out = self.deconv1(out, x1)
        out = self.output(out)
        return out



